import torch
import numpy as np
import matplotlib.pyplot as plt
import torch.nn as nn
import torch.optim

torch.manual_seed(10)

n = 10000
mean_value = 2.5  # mean value of the distribution
bias = 1.3

n_data = torch.ones(n, 2)
x0 = torch.normal(mean_value * n_data, 1.0) + bias
y0 = torch.zeros(n)

x1 = torch.normal(-mean_value * n_data, 1.0) + bias
y1 = torch.ones(n)

train_x_0 = torch.cat((x0, x1), 0)
train_y_0 = torch.cat((y0, y1), 0)

import pandas as pd

A = pd.read_csv('/Users/lizheng/PycharmProjects/tianchi/data/watermelon3.csv', sep=' ')
train_x_0 = torch.Tensor(A.iloc[:, :2].values)
y_l = A.iloc[:, -1].values
train_y_0 = torch.Tensor([1 if y_ == 1 else 0 for y_ in y_l])
print(train_x_0.shape, train_y_0.shape)

np.random.seed(10)
ind = np.random.permutation(range(train_y_0.shape[0]))
train_x, train_y = train_x_0[ind], train_y_0[ind]
n = train_y.shape[0]
ntrn = int(n * 0.8)
x_trn, y_trn, x_val, y_val = train_x[:ntrn], train_y[:ntrn], train_x[ntrn:], train_y[ntrn:]


class LogReg(nn.Module):
    def __init__(self):
        super(LogReg, self).__init__()
        self.features1 = nn.Linear(2, 1, bias=True)
        # self.features2 = nn.Linear(5, 1)
        self.act_func = nn.Sigmoid()

    def forward(self, x):
        x = self.features1(x)
        # x = self.features2(x)
        x = self.act_func(x)
        return x


def eval_loss_acc(trn, y, model):
    y_ = model(trn)
    loss = criterion(y_.squeeze(), y)
    mask = y_.ge(0.5).float().squeeze()
    correct = (mask == y).sum()
    acc = correct.item() / y_.shape[0]
    return loss, acc


model = LogReg()
# model.features1.reset_parameters()
print(model.features1.weight)

# w, b = [[-1000.00, -6.35447905]], [0.000]
# model.features1.weight.data = torch.tensor(w)
# model.features1.bias.data = torch.tensor(b)

criterion = nn.MSELoss()
lr = 0.001
threshold = 0.1
optimizer = torch.optim.SGD(model.parameters(), lr=lr)
# optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

xl, yl, Acc = [], [], []
xl_val, yl_val, Acc_val = [], [], []
batch_sz = 50000
N = y_trn.shape[0]
last_loss = 0
for step in range(300):
    si = step * batch_sz % N
    ei = (step + 1) * batch_sz % N
    if ei > si:
        xtrn_step = x_trn[si:ei]
        y_trn_step = y_trn[si:ei]
    else:
        xtrn_step = torch.cat([x_trn[si:], x_trn[:ei]], dim=0)
        y_trn_step = torch.cat([y_trn[si:], y_trn[:ei]], dim=0)
    if y_trn_step.shape[0] != batch_sz:
        print('-' * 100)
    y_pred = model(xtrn_step)
    loss1 = criterion(y_pred.squeeze(), y_trn_step)
    l1_add_1 = (model.features1.weight - torch.tensor([0, 0])).norm()
    l1_add_2 = (model.features1.bias - torch.tensor([0])).norm()
    l1_add_loss = l1_add_1 + l1_add_2
    l1_add_loss = l1_add_1 + 0
    # l1_add_loss = 0
    loss = loss1 + l1_add_loss
    loss.backward()

    optimizer.step()
    if step % 1 == 0:
        # 计算准确率
        mask = y_pred.ge(0.5).float().squeeze()
        correct = (mask == y_trn_step).sum()
        acc = correct.item() / y_trn_step.shape[0]
        xl.append(step)
        yl.append(loss1.item())
        Acc.append(acc)
        if abs(last_loss - loss1.item()) < threshold:
            last_lr = lr
            lr = last_lr * 0.01
            threshold = threshold * 0.5
            optimizer = torch.optim.SGD(model.parameters(), lr=lr)
            print(step, '-' * 50, lr)
            # break
        loss_val, acc_val = eval_loss_acc(x_val, y_val, model)
        xl_val.append(step)
        yl_val.append(loss_val.item())
        Acc_val.append(acc_val)
        print(step, model.features1.weight.detach().numpy(), model.features1.bias.detach().numpy())
        print('epoch: ', step, 'loss: ', loss.item(), 'acc: ', acc)
        print('epoch-val: ', step, 'loss: ', loss_val.item(), 'acc: ', acc_val)

from matplotlib import pyplot as plt

plt.plot(xl, yl, xl, Acc, xl_val, yl_val, xl_val, Acc_val)
plt.legend(['trn-loss', 'trn-acc', 'val-loss', 'val-acc'])
plt.show()

# w, b = [[-6.89352917, -6.35447905]], [16.44745518]
# model.features1.weight.data = torch.tensor(w)
# model.features1.bias.data = torch.tensor(b)
# a = model(torch.tensor([[1.00, 0.00],
#                         [1.00, 1.00],
#                         [0.00, 1.00],
#                         [4.7804, 4.0977]
#                         ]))
# print(a)
